{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Static forward free-boundary equilibrium calculations\n",
    "\n",
    "---\n",
    "\n",
    "This example notebook shows how to use FreeGSNKE to solve **static forward** free-boundary Grad-Shafranov (GS) problems.\n",
    "\n",
    "In the **forward** solve mode we solve for the plasma equilibrium using user-defined active poloidal field coil currents, passive structure currents, and plasma current density profiles. \n",
    "\n",
    "Below, we illustrate how to use the solver for both diverted and limited plasma configurations in a **MAST-U-like tokamak** using stored pickle files containing the machine description. These machine description files partially come from the FreeGS repository and are not an exact replica of MAST-U. \n",
    "\n",
    "##### Note:\n",
    "It is recommended to go through the inverse solver notebook before this one as we omit many of the commonly shared details!\n",
    "\n",
    "We'll now go through the steps required to solve the **forward** problem in FreeGSNKE. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import some packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the machine object\n",
    "\n",
    "First, we build the machine object from previously created pickle files in the \"machine_configs/MAST-U\" directory. \n",
    "\n",
    "FreeGSNKE requires the following paths in order to build the machine:\n",
    "- `active_coils_path`\n",
    "- `passive_coils_path`\n",
    "- `limiter_path`\n",
    "- `wall_path`\n",
    "- `magnetic_probe_path` (not required here)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate an equilibrium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import equilibrium_update\n",
    "\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,      # provide tokamak object\n",
    "    Rmin=0.1, Rmax=2.0,   # radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # vertical range\n",
    "    nx=65,                # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    # psi=plasma_psi\n",
    ")  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate a profile object\n",
    "\n",
    "See inverse solver notebook for a list of some of the different profile objects available in FreeGSNKE. Later on, we will try some of these out. \n",
    "\n",
    "Here, we will use the `ConstrainPaxisIp` profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise the profiles\n",
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,        # equilibrium object\n",
    "    paxis=8e3,    # profile object\n",
    "    Ip=6e5,       # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=1.8,  # profile function parameter\n",
    "    alpha_n=1.2   # profile function parameter\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the static nonlinear solver\n",
    "\n",
    "We can now load FreeGSNKE's Grad-Shafranov static solver. The equilibrium is used to inform the solver of the computational domain and of the tokamak properties. The solver below can be used for both inverse and forward solve modes.\n",
    "\n",
    "Note: It's not necessary to instantiate a new solver when aiming to use it on new or different equilibria, as long as the integration domain, mesh grid, and tokamak are consistent across solves. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the coil currents\n",
    "\n",
    "As mentioned before, during a forward solve, we use fixed coil currents (as well as given profile functions/parameters) as inputs to solve for the equilibrium. To do this, we can use the set of currents we identified within the inverse solve notebook. Note that the passive structure currents are zero. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the coil currents\n",
    "import pickle\n",
    "with open('simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n",
    "    currents_dict = pickle.load(f)\n",
    "    \n",
    "# assign currents to the eq object\n",
    "for key in currents_dict.keys():\n",
    "    eq.tokamak.set_coil_current(coil_label=key, current_value=currents_dict[key])\n",
    "    \n",
    "eq.tokamak.getCurrents()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The forward solve\n",
    "\n",
    "The syntax of a forward solve is identical to that of an inverse call (i.e. calling `GSStaticSolver.solve()`), however the nonlinear solver is **not** provided with a `constrain` object (i.e. we set `constrain=None`). \n",
    "\n",
    "Therefore, coil current values are not modified during the solve but instead the solver uses them as inputs to calculate the equilibrium. \n",
    "\n",
    "FreeGSNKE uses a Newton-Krylov (NK) method to solve the static forward problem stated in the inverse solve notebook (which we re-write below as a root-problem, boundary condition omitted here):\n",
    "\n",
    "$$ \\Delta^* \\psi + \\mu_0 R J_{\\phi}(\\psi,R,Z) = 0, \\qquad (R,Z) \\in \\Omega. $$\n",
    "\n",
    "Given the coil flux $\\psi_c$ is known prior to solving, we only require an initial guess for the plasma flux $\\psi_p$. This is generated automatically in FreeGSNKE and scaled automatically according to the size of the coil currents and/or plasma current. If a good initial guess is known, it can be provided to the solver in the equilbirium object above via the `psi` option. \n",
    "\n",
    "The NK method helps mitigate the numerical instability problems associated with Picard-based iterations and enables considerably more restrictive tolerance requests. The stopping criterion is defined as\n",
    "\n",
    "$$ \\frac{\\text{max} | \\psi^{(n+1)} - \\psi^{(n)} |}{\\text{max} \\ \\psi^{(n)} - \\text{min} \\ \\psi^{(n)}} < \\varepsilon, $$\n",
    "\n",
    "where $n$ is the iteration number. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# call the solver\n",
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9,\n",
    "                     verbose=True, # print output\n",
    "                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What we have done here is improve on the equilibrium found during the inverse solve by taking the coil currents (found during the inverse solve) and the same profile parameters and feeding them into the forward solver. \n",
    "\n",
    "We do this because it is often difficult to achieve low relative tolerances in _inverse_ solve calls (for example, the above was set at a loose target_relative_tolerance=1e-6) and so the strategy of using a forward solve after an inverse one is useful to obtain better (more converged) equilibria at stricter tolerances.\n",
    "\n",
    "As an additional example, below we manually vary some of the coil currents, perform new forward solves and compare the resulting equilibria.\n",
    "Note that the manual current changes cause the equilibria to transition from a diverted to a limiter configuration (this is handled through FreeGS4E)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "# copy the original eq object (for the new forward solves with modified currents)\n",
    "eq_forward_1 = deepcopy(eq)\n",
    "eq_forward_2 = deepcopy(eq)\n",
    "\n",
    "# modify the P4 current and solve\n",
    "eq_forward_1.tokamak.set_coil_current('P4', 1.5*eq.tokamak['P4'].current)\n",
    "GSStaticSolver.solve(eq=eq_forward_1, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "# modify the P4 current (even more) and solve\n",
    "eq_forward_2.tokamak.set_coil_current('P4', 1.5**2 * eq.tokamak['P4'].current)\n",
    "GSStaticSolver.solve(eq=eq_forward_2, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "\n",
    "# plot the resulting equilbria \n",
    "fig1, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 8), dpi=80)\n",
    "\n",
    "# original \n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1,show=False)\n",
    "eq.tokamak.plot(axis=ax1,show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "\n",
    "# modified 1\n",
    "ax2.grid(True, which='both')\n",
    "eq_forward_1.plot(axis=ax2,show=False)\n",
    "eq_forward_1.tokamak.plot(axis=ax2,show=False)\n",
    "ax2.set_xlim(0.1, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "\n",
    "# modified 2\n",
    "ax3.grid(True, which='both')\n",
    "eq_forward_2.plot(axis=ax3,show=False)\n",
    "eq_forward_2.tokamak.plot(axis=ax3,show=False)\n",
    "ax3.set_xlim(0.1, 2.15)\n",
    "ax3.set_ylim(-2.25, 2.25)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternative profile functions\n",
    "Here, we will illustrate how to use some of the other available profile objects.\n",
    "\n",
    "We'll start with the `ConstrainBetapIp` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import ConstrainBetapIp\n",
    "\n",
    "profiles_beta = ConstrainBetapIp(\n",
    "    eq=eq,\n",
    "    betap=0.05,\n",
    "    Ip=6e5,\n",
    "    fvac=0.5,\n",
    "    alpha_m=1.8,\n",
    "    alpha_n=1.2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use it directly in a new solve (with the coil currents found by the inverse solve)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instatiate new equilibrium object\n",
    "eq_beta = deepcopy(eq)\n",
    "\n",
    "# call solver with new profile object\n",
    "GSStaticSolver.solve(eq=eq_beta, \n",
    "                     profiles=profiles_beta, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "\n",
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq_beta.plot(axis=ax1, show=False)\n",
    "eq_beta.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use the `Fiesta_Topeol` profile (see equation 13 in L. L. Lao et al (1985) Nucl. Fusion 25 1611). This has the same parameterisation as the previous two profiles except that we can now specify the  `Beta0` parameter directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import Fiesta_Topeol\n",
    "\n",
    "profiles_topeol = Fiesta_Topeol(\n",
    "    eq=eq,        # equilibrium object\n",
    "    Beta0=0.3665, # beta0 parameter\n",
    "    Ip=6e5,       # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=2,    # profile function parameter\n",
    "    alpha_n=1     # profile function parameter\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instatiate new equilibrium object\n",
    "eq_topeol = deepcopy(eq)\n",
    "\n",
    "# call solver with new profile object\n",
    "GSStaticSolver.solve(eq=eq_topeol, \n",
    "                     profiles=profiles_topeol, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "\n",
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq_topeol.plot(axis=ax1, show=False)\n",
    "eq_topeol.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now test the `Lao85` profile (see equations 4 and 5 in L. L. Lao et al (1985) Nucl. Fusion 25 1611).\n",
    "\n",
    "These profiles are parametrised as:\n",
    "\n",
    " $$J_{p}(\\psi, R, Z) = \\lambda\\big[ \\frac{R}{R_{0}} p'(\\tilde{\\psi}) + \\frac{R_0}{R} \\frac{1}{\\mu_0} F F'(\\tilde{\\psi}) \\big] \\quad (R,Z) \\in \\Omega_p, $$\n",
    "where the pressure and toroidal current profiles are given by\n",
    "$$ p'(\\tilde{\\psi}) = \\sum_{i=0}^{n_p} \\alpha_i \\tilde{\\psi}^i - \\hat{\\alpha} \\tilde{\\psi}^{n_p + 1} \\sum_{i=0}^{n_p} \\alpha_i$$\n",
    "and\n",
    "$$ F F'(\\tilde{\\psi}) = \\sum_{i=0}^{n_F} \\beta_i \\tilde{\\psi}^i - \\hat{\\beta} \\tilde{\\psi}^{n_F + 1} \\sum_{i=0}^{n_F} \\beta_i.$$ \n",
    "\n",
    "The required parameters are:\n",
    "- `Ip` (total plasma current).\n",
    "- `fvac` ($rB_{tor}$,  vacuum toroidal field strength).\n",
    "- `alpha` (array of alpha coefficients).\n",
    "- `beta` (array of beta coefficients).\n",
    "- `alpha_logic`, and `beta_logic` (Booleans that correspond to $\\hat{\\alpha}$ and $\\hat{\\beta}$ above, sets boundary condition for plasma current at plasma boundary).\n",
    "- `Ip_logic` (if False, `Ip` is not used, if True, `Ip` is used to normalise $J_p$ and find $\\lambda$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the Lao85 profile to set up the same identical equilibrium as generated by the Topeol profile. \n",
    "\n",
    "In the Topeol profiles, we have $\\alpha_m = 2$ and $\\alpha_n = 1$ which means both $p'$ and $FF'$ are proportional to $1 - \\psi_n^2$ (with $\\hat{\\alpha} = \\hat{\\beta} = 1$. This corresponds to (the vectors) $\\alpha, \\beta \\propto (1, 0 , -1)$ - check this. \n",
    "\n",
    "We also need to take into account the scalings: alpha over beta is mu0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from freegsnke.jtor_update import Lao85\n",
    "from freegs4e.gradshafranov import mu0   # permeability\n",
    "\n",
    "alpha = np.array([1,0,-1])\n",
    "beta = (1 - profiles_topeol.Beta0)/profiles_topeol.Beta0 * alpha * mu0\n",
    "\n",
    "profiles_lao = Lao85(\n",
    "    eq=eq,\n",
    "    Ip=6e5,\n",
    "    fvac=0.5,\n",
    "    alpha=alpha,\n",
    "    beta=beta,\n",
    "    alpha_logic=False,\n",
    "    beta_logic=False,\n",
    "    Ip_logic=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that above we're providing as input the full list of alpha and beta coefficients and therefore setting both logic inputs to False.\n",
    "The following is entirely equivalent to the above:\n",
    "\n",
    "```python\n",
    "alpha = np.array([1,0])\n",
    "beta = (1 - profiles_topeol.Beta0)/profiles_topeol.Beta0 * alpha * mu0\n",
    "\n",
    "profiles_lao = Lao85(\n",
    "    eq=eq_forward,\n",
    "    Ip=6e5,\n",
    "    fvac=0.5,\n",
    "    alpha=alpha,\n",
    "    beta=beta,\n",
    "    alpha_logic=True,\n",
    "    beta_logic=True,\n",
    "    Ip_logic=True,\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instatiate new equilibrium object\n",
    "eq_lao = deepcopy(eq)\n",
    "\n",
    "# call solver with new profile object\n",
    "GSStaticSolver.solve(eq=eq_lao, \n",
    "                     profiles=profiles_lao, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "\n",
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq_lao.plot(axis=ax1, show=False)\n",
    "eq_lao.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following illustrates that the two profile functions indeed generate the same current distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "laoj = profiles_lao.Jtor(R=eq.R, Z=eq.Z, psi=eq.psi())\n",
    "topj = profiles_topeol.Jtor(R=eq.R, Z=eq.Z, psi=eq.psi())\n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(5, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "plt.contourf(eq.R, eq.Z, (laoj-topj))\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "plt.plot(eq.tokamak.wall.R, eq.tokamak.wall.Z, 'k', 1.2)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "FreeGSNKE also enables one to specify a set of Topeol profile parameters that best fit a set of Lao85 parameters (using the Lao_parameters method)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha, beta = profiles_topeol.Lao_parameters(n_alpha=2, n_beta=2, alpha_logic=True, beta_logic=True)\n",
    "print(f\"Original alpha's = {profiles_lao.alpha[0:2]} vs. Fitted from Topeol = {alpha}.\")\n",
    "print(f\"Original beta's = {profiles_lao.beta[0:2]} vs. Fitted from Topeol = {beta}.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles_lao_fit = Lao85(\n",
    "    eq=eq,\n",
    "    Ip=6e5,\n",
    "    fvac=0.5,\n",
    "    alpha=alpha,\n",
    "    beta=beta,\n",
    "    alpha_logic=True,\n",
    "    beta_logic=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "laoj = profiles_lao_fit.Jtor(R=eq.R, Z=eq.Z, psi=eq.psi())\n",
    "topj = profiles_topeol.Jtor(R=eq.R, Z=eq.Z, psi=eq.psi())\n",
    "\n",
    "\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(5, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "plt.contourf(eq.R, eq.Z, (laoj-topj))\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "plt.plot(eq.tokamak.wall.R, eq.tokamak.wall.Z, 'k', 1.2)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reverse is also possible, using the the Topeol_parameters method in the Lao85 profile object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_m, alpha_n, beta_0 = profiles_lao_fit.Topeol_parameters()\n",
    "\n",
    "print(f\"Original alpha_m = {profiles_topeol.alpha_m} vs. Fitted from Lao85 = {alpha_m}.\")\n",
    "print(f\"Original alpha_n = {profiles_topeol.alpha_n} vs. Fitted from Lao85 = {alpha_n}.\")\n",
    "print(f\"Original beta_0 = {profiles_topeol.Beta0} vs. Fitted from Lao85 = {beta_0}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forward solve: limiter plasma\n",
    "\n",
    "Here we use the saved limiter plasma currents from the prior notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise the equilibrium\n",
    "eq_limiter = deepcopy(eq)\n",
    "\n",
    "\n",
    "# initialise the profiles\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq_limiter,        # equilibrium object\n",
    "    paxis=6e3,    # profile object\n",
    "    Ip=4e5,       # plasma current\n",
    "    fvac=0.5,     # fvac = rB_{tor}\n",
    "    alpha_m=1.8,  # profile function parameter\n",
    "    alpha_n=1.2   # profile function parameter\n",
    ")\n",
    "\n",
    "# load the nonlinear solver\n",
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq_limiter)    \n",
    "\n",
    "# set the coil currents\n",
    "import pickle\n",
    "with open('simple_limited_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq_limiter.tokamak.set_coil_current(coil_label=key, current_value=current_values[key])\n",
    "\n",
    "# carry out the foward solve to find the equilibrium\n",
    "GSStaticSolver.solve(eq=eq_limiter, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "\n",
    "# plot the resulting equilbria \n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq_limiter.plot(axis=ax1, show=False)\n",
    "eq_limiter.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "freegsnke4e",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}